#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SM2Algorithm的泄漏模型
定义SM2各个操作的Power泄漏特征
"""

import numpy as np
from typing import Dict, Any, List, Tuple, Optional


class SM2LeakageModels:
    """SM2泄漏模型集合"""
    
    # SM2推荐曲线Parameters
    SM2_P = 0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF
    SM2_A = 0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC
    SM2_B = 0x28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93
    SM2_N = 0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123
    SM2_GX = 0x32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7
    SM2_GY = 0xBC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0
    
    @staticmethod
    def hamming_weight(value: int) -> int:
        """Calculate汉明权重"""
        return bin(value).count('1')
    
    @staticmethod
    def hamming_distance(a: int, b: int) -> int:
        """CalculateHamming distance"""
        return SM2LeakageModels.hamming_weight(a ^ b)
    
    @classmethod
    def sm2_point_addition_model(cls, x1: Optional[int], y1: Optional[int], 
                                 x2: Optional[int], y2: Optional[int], 
                                 p: int) -> Tuple[Optional[int], Optional[int]]:
        """
        SM2Elliptic Curve点加法泄漏模型
        
        Args:
            x1, y1: 点P1的坐标（None表示No穷远点）
            x2, y2: 点P2的坐标（None表示No穷远点）
            p: 素数模数
            
        Returns:
            点P1 + P2的坐标（None, None表示No穷远点）
        """
        # ProcessNo穷远点
        if x1 is None or y1 is None:
            return (x2, y2)
        if x2 is None or y2 is None:
            return (x1, y1)
        
        # CheckYesNoFor相同点Or互For逆元
        if x1 == x2:
            if y1 == y2:
                # 相同点 -> 点倍乘
                return cls.sm2_point_doubling_model(x1, y1, p)
            else:
                # 互For逆元 -> No穷远点
                return (None, None)
        
        # Calculate斜率
        s = ((y2 - y1) * pow(x2 - x1, -1, p)) % p
        
        # Calculate新点坐标
        x3 = (s * s - x1 - x2) % p
        y3 = (s * (x1 - x3) - y1) % p
        
        return (x3, y3)
    
    @classmethod
    def sm2_point_doubling_model(cls, x: int, y: int, p: int, a: int = None) -> Tuple[int, int]:
        """
        SM2Elliptic Curve点倍乘泄漏模型
        
        Args:
            x, y: 点P的坐标
            p: 素数模数
            a: 曲线Parametersa（DefaultUseSM2Parameters）
            
        Returns:
            点2P的坐标
        """
        if a is None:
            a = cls.SM2_A
        
        # Calculate斜率 s = (3x^2 + a) / (2y)
        numerator = (3 * x * x + a) % p
        denominator = (2 * y) % p
        s = (numerator * pow(denominator, -1, p)) % p
        
        # Calculate新点坐标
        x3 = (s * s - 2 * x) % p
        y3 = (s * (x - x3) - y) % p
        
        return (x3, y3)
    
    @classmethod
    def sm2_scalar_multiplication_model(cls, k: int, x: int, y: int, p: int) -> Tuple[Optional[int], Optional[int]]:
        """
        SM2Scalar乘法泄漏模型（Double-and-AddAlgorithm）
        
        Args:
            k: Scalar（私钥）
            x, y: 基点坐标
            p: 素数模数
            
        Returns:
            点kP的坐标（None, None表示No穷远点）
        """
        if k == 0:
            return (None, None)  # No穷远点
        
        if k < 0:
            raise ValueError("Scalark必须For非负数")
        
        # Double-and-AddAlgorithm（From高位To低位）
        k_bin = bin(k)[2:]  # 去掉'0b'
        
        # InitializeResultFor基点P
        result_x, result_y = x, y
        
        # From第二位Start（最高位已经Process）
        for bit in k_bin[1:]:
            # 总YesExecute点倍乘
            result_x, result_y = cls.sm2_point_doubling_model(result_x, result_y, p)
            
            # IfCurrent位For1，Execute点加法
            if bit == '1':
                result_x, result_y = cls.sm2_point_addition_model(result_x, result_y, x, y, p)
        
        return (result_x, result_y)
    
    @classmethod
    def sm2_sign_model(cls, message_hash: int, private_key: int, nonce: int, 
                      base_x: int, base_y: int, p: int, n: int) -> Tuple[int, int]:
        """
        SM2签名泄漏模型
        
        Args:
            message_hash: 消息哈希
            private_key: 私钥
            nonce: 随机数k
            base_x, base_y: 基点坐标
            p: 曲线素数
            n: 阶
            
        Returns:
            签名(r, s)
        """
        # Calculate r = (e + x1) mod n，其中(x1, y1) = k * G
        kG_x, _ = cls.sm2_scalar_multiplication_model(nonce, base_x, base_y, p)
        r = (message_hash + kG_x) % n
        
        if r == 0 or (r + nonce) % n == 0:
            return (0, 0)  # Require重新选择k
        
        # Calculate s = (1 + private_key)^(-1) * (k - r * private_key) mod n
        numerator = (nonce - r * private_key) % n
        denominator = (1 + private_key) % n
        s = (numerator * pow(denominator, -1, n)) % n
        
        return (r, s)
    
    @classmethod
    def sm2_encrypt_model(cls, message: int, public_key_x: int, public_key_y: int,
                         nonce: int, base_x: int, base_y: int, p: int) -> Tuple[int, int, int]:
        """
        SM2加密泄漏模型（简化版）
        
        Args:
            message: 消息
            public_key_x, public_key_y: 公钥坐标
            nonce: 随机数k
            base_x, base_y: 基点坐标
            p: 素数模数
            
        Returns:
            Ciphertext(C1_x, C1_y, C2)
        """
        # C1 = k * G
        C1_x, C1_y = cls.sm2_scalar_multiplication_model(nonce, base_x, base_y, p)
        
        # Calculate k * P_B（公钥点）
        kPB_x, kPB_y = cls.sm2_scalar_multiplication_model(nonce, public_key_x, public_key_y, p)
        
        # C2 = message XOR KDF(kPB_x || kPB_y)（简化For直接XOR）
        C2 = message ^ (kPB_x & 0xFFFFFFFF)  # 简化的Key派生
        
        return (C1_x, C1_y, C2)
    
    @classmethod
    def sm2_key_exchange_model(cls, private_key: int, temp_private_key: int,
                             peer_public_x: int, peer_public_y: int,
                             peer_temp_x: int, peer_temp_y: int,
                             base_x: int, base_y: int, p: int, n: int) -> int:
        """
        SM2Key交换泄漏模型（简化版）
        
        Args:
            private_key: 长期私钥
            temp_private_key: 临时私钥
            peer_public_x, peer_public_y: 对方公钥
            peer_temp_x, peer_temp_y: 对方临时公钥
            base_x, base_y: 基点坐标
            p: 素数模数
            n: 阶
            
        Returns:
            共享Key（简化Forx坐标）
        """
        # Calculate t = (private_key + temp_private_key * x) mod n
        # 这里简化Process
        t = (private_key + temp_private_key) % n
        
        # Calculate共享点 [t](P_B + [x]R_B)
        # 简化For直接Calculate t * peer_public
        shared_x, _ = cls.sm2_scalar_multiplication_model(t, peer_public_x, peer_public_y, p)
        
        return shared_x
    
    @classmethod
    def sm3_compression_model(cls, message_block: List[int]) -> List[int]:
        """
        SM3哈希压缩Function泄漏模型（简化版）
        
        Args:
            message_block: 消息块（16个32位字）
            
        Returns:
            压缩后的哈希值（8个32位字）
        """
        # 简化的SM3压缩Function
        # 实际SM3更复杂，这里ProvideBasic框架
        
        # Initialize
        A, B, C, D, E, F, G, H = [
            0x7380166f, 0x4914b2b9, 0x172442d7, 0xda8a0600,
            0xa96f30bc, 0x163138aa, 0xe38dee4d, 0xb0fb0e4e
        ]
        
        # 简化的压缩过程
        for i in range(64):
            # 简化的轮Function
            T = 0x79cc4519 if i < 16 else 0x7a879d8a
            SS1 = ((A << 12) | (A >> 20)) + E + ((T << (i % 32)) | (T >> (32 - i % 32)))
            SS1 = ((SS1 << 7) | (SS1 >> 25)) & 0xffffffff
            SS2 = SS1 ^ ((A << 12) | (A >> 20))
            
            # UpdateStatus
            A, B, C, D, E, F, G, H = (
                SS2, A, (B << 9) | (B >> 23), C, SS1, E, (F << 19) | (F >> 13), G
            )
            A, B, C, D, E, F, G, H = [x & 0xffffffff for x in [A, B, C, D, E, F, G, H]]
        
        return [A, B, C, D, E, F, G, H]
    
    @classmethod
    def get_intermediate_functions(cls) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return {
            'point_addition': cls.sm2_point_addition_model,
            'point_doubling': cls.sm2_point_doubling_model,
            'scalar_multiplication': cls.sm2_scalar_multiplication_model,
            'sm2_sign': cls.sm2_sign_model,
            'sm2_encrypt': cls.sm2_encrypt_model,
            'sm2_key_exchange': cls.sm2_key_exchange_model,
            'sm3_compression': cls.sm3_compression_model,
        }
    
    @classmethod
    def get_leakage_models(cls) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return {
            'value': lambda x: x,
            'hamming_weight': cls.hamming_weight,
            'hamming_distance': cls.hamming_distance,
            'bit_value': lambda x, bit_pos: (x >> bit_pos) & 1,
        }
    
    @classmethod
    def get_target_operations(cls) -> List[str]:
        """获取可Attack的目标操作"""
        return [
            'scalar_multiplication',
            'point_addition',
            'point_doubling',
            'sm2_signing',
            'sm2_encryption',
            'sm2_key_exchange',
            'sm3_hashing',
            'field_multiplication',
            'field_inversion',
            'modular_reduction'
        ]
    
    @classmethod
    def get_operation_to_intermediate_mapping(cls) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return {
            'scalar_multiplication': 'scalar_multiplication',
            'point_addition': 'point_addition',
            'point_doubling': 'point_doubling',
            'sm2_signing': 'sm2_sign',
            'sm2_encryption': 'sm2_encrypt',
            'sm2_key_exchange': 'sm2_key_exchange',
            'sm3_hashing': 'sm3_compression',
            'field_multiplication': 'point_addition',
            'field_inversion': 'point_addition',
            'modular_reduction': 'scalar_multiplication',
        }
    
